from langchain_core.documents import Document
from tqdm import tqdm
from weaviate.collections.classes.filters import Filter
from weaviate.util import generate_uuid5
from weaviate.exceptions import WeaviateBaseError

from utils.tools import create_vector
from .database import db


# 创建集合
def create_collection(collection_name: str):
    client = db.get_client()
    client.collections.create(name=collection_name)

# 删除集合
def delete_collection(collection_name: str):
    client = db.get_client()
    client.collections.delete(name=collection_name)

# 创建数据
def create_object(collection_name: str, properties: dict, class_name: str = None, vector: list = None):
    client = db.get_client()
    try:
        jeopardy = client.collections.get(collection_name)
        if class_name is None:
            jeopardy.data.insert(properties=properties, vector=vector)
        else:
            jeopardy.data.insert(properties=properties, uuid=generate_uuid5(class_name), vector=vector)
        return {"status": "created"}
    except WeaviateBaseError as e:
        client.close()
        raise Exception(f"Failed to create object: {e}")


# 批量导入
def create_objects(collection_name: str, objects: list[Document]):
    client = db.get_client()
    try:
        jeopardy = client.collections.get(collection_name)
        with jeopardy.batch.dynamic() as batch:
            for i, data_row in tqdm(enumerate(objects)):
                properties = data_row.metadata
                properties["text"] = data_row.page_content
                vector = create_vector(data_row.page_content)
                batch.add_object(
                    properties=properties,
                    vector=vector
                )
        return {"status": "created"}
    except WeaviateBaseError as e:
        client.close()
        raise Exception(f"Failed to create object: {e}")


# 读取数据
def get_object(collection_name: str, class_name: str):
    client = db.get_client()
    try:
        jeopardy = client.collections.get(collection_name)
        response = jeopardy.query.fetch_object_by_id(generate_uuid5(class_name))
        return response
    except WeaviateBaseError as e:
        client.close()
        raise Exception(f"Failed to retrieve object: {e}")


# 更新数据
def update_object(collection_name: str, class_name: str, properties: dict):
    client = db.get_client()
    try:
        jeopardy = client.collections.get(collection_name)
        jeopardy.data.update(properties=properties, uuid=generate_uuid5(class_name))
        return {"status": "updated"}
    except WeaviateBaseError as e:
        client.close()
        raise Exception(f"Failed to update object: {e}")


# 删除数据
def delete_object(collection_name: str, class_name: str):
    client = db.get_client()
    try:
        jeopardy = client.collections.get(collection_name)
        jeopardy.data.delete_by_id(generate_uuid5(class_name))
        return {"status": "deleted"}
    except WeaviateBaseError as e:
        client.close()
        raise Exception(f"Failed to delete object: {e}")

# 根据字段删除数据
def delete_object_by_field(collection_name: str, field_name: str, field_value: str):
    client = db.get_client()
    try:
        jeopardy = client.collections.get(collection_name)
        jeopardy.data.delete_many(
            where=Filter.by_property(field_name).equal(field_value)
        )
        return {"status": "deleted"}
    except WeaviateBaseError as e:
        print(f"Failed to delete object: {e}")